import re

# 读取test_env.py
with open('test_env.py', 'r') as f:
    content = f.read()

# 修复gym导入
content = content.replace('import gym', 'import gymnasium as gym')

# 修复tensor维度问题
content = re.sub(
    r'g_action_space = gym\.spaces\.Box\(low=0\.0, high=1\.0, shape=\(2,\), dtype=np\.float32\)',
    'g_action_space = gym.spaces.Box(low=0.0, high=1.0, shape=(2,), dtype=np.float32)',
    content
)

# 在Global Policy测试部分添加正确的参数
content = re.sub(
    r'g_policy = RL_Policy\(g_observation_space\.shape, g_action_space,\s*base_kwargs=\{\'recurrent\': False, \'hidden_size\': 256\}\)\.to\(args\.device\)',
    '''g_policy = RL_Policy(g_observation_space.shape, g_action_space,
                           base_kwargs={'recurrent': False, 'hidden_size': 256,
                                      'downscaling': 1}).to(args.device)''',
    content
)

# 写回文件
with open('test_env.py', 'w') as f:
    f.write(content)

print("✅ test_env.py已修复")